package com.webchat.aigc.llm;

import com.alibaba.dashscope.embeddings.TextEmbedding;
import com.alibaba.dashscope.embeddings.TextEmbeddingParam;
import com.alibaba.dashscope.embeddings.TextEmbeddingResult;
import com.webchat.aigc.config.properties.AliBabaEmbeddingPropertiesConfig;
import org.apache.commons.lang3.ArrayUtils;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Service;

import java.util.Arrays;
import java.util.List;


@Service
public class AlibabaEmbeddingModel extends AbstractEmbeddingModel {

    @Autowired
    private AliBabaEmbeddingPropertiesConfig embeddingPropertiesConfig;


    @Override
    public float[] embed(String text) throws Exception {
        TextEmbeddingParam param = TextEmbeddingParam
                .builder()
                .model(TextEmbedding.Models.TEXT_EMBEDDING_V3)
                .apiKey(embeddingPropertiesConfig.getApiKey())
                .text(text)
                .build();
        TextEmbedding textEmbedding = new TextEmbedding();
        TextEmbeddingResult result = textEmbedding.call(param);
        List<Double> embeddings = result.getOutput().getEmbeddings().get(0).getEmbedding();
        List<Float> fEmbeddings = embeddings.stream().map(embed -> embed.floatValue()).toList();
        return ArrayUtils.toPrimitive(fEmbeddings.toArray(new Float[0]));
    }

    @Override
    public List<float[]> embed(List<String> texts) throws Exception {
        TextEmbeddingParam param = TextEmbeddingParam
                .builder()
                .apiKey(embeddingPropertiesConfig.getApiKey())
                .model(TextEmbedding.Models.TEXT_EMBEDDING_V3)
                .texts(texts).build();
        TextEmbedding textEmbedding = new TextEmbedding();
        TextEmbeddingResult result = textEmbedding.call(param);
        return result.getOutput().getEmbeddings().stream().map(embed -> {
            List<Double> embeddings = embed.getEmbedding();
            List<Float> fEmbeddings = embeddings.stream().map(textEmbed ->
                                                                    textEmbed.floatValue()).toList();
            return ArrayUtils.toPrimitive(fEmbeddings.toArray(new Float[0]));
        }).toList();
    }
}
